import sys
import serial.tools.list_ports
import os
import re
import time
from PyQt5.QtWidgets import (QApplication, QTableWidget, QTableWidgetItem,
                             QVBoxLayout, QWidget, QProgressBar, QPushButton,
                             QCheckBox, QMessageBox, QHBoxLayout, QFileDialog,
                             QDialog, QPlainTextEdit, QAbstractItemView, QFrame,
                             QHeaderView, QLabel, QStyleFactory)
from PyQt5.QtCore import QThread, pyqtSignal, QObject, QTimer, QProcess, Qt, pyqtSlot
from PyQt5.QtGui import QColor, QBrush, QFont, QIcon, QPixmap
from threading import Lock # 导入线程锁
import openpyxl

# ==============================================================================
# 1. ATCommunicator 类 
# ==============================================================================
class ATCommunicator:
    """封装所有AT命令串口通信的逻辑。"""
    def __init__(self, port, baudrate=115200, timeout=2):
        self.port = port
        self.baudrate = baudrate
        self.timeout = timeout
        self.ser = None

    def _open(self):
        try:
            self.ser = serial.Serial(
                port=self.port, 
                baudrate=self.baudrate, 
                timeout=self.timeout,
                xonxoff=False,
                rtscts=False,
                dsrdtr=False
            )
            time.sleep(0.1)
            return True, ""
        except serial.SerialException as e:
            return False, f"打开端口 {self.port} 失败: {e}"

    def _close(self):
        if self.ser and self.ser.is_open:
            self.ser.close()

    def send_command(self, command, response_timeout=5):
        is_open, msg = self._open()
        if not is_open:
            return False, [msg]
        try:
            self.ser.reset_input_buffer()
            self.ser.reset_output_buffer()
            time.sleep(0.2)
            self.ser.write((command + '\r\n').encode('ascii'))
            response_lines = []
            start_time = time.time()
            while time.time() - start_time < response_timeout:
                line = self.ser.readline().decode('utf-8', errors='ignore').strip()
                if line:
                    response_lines.append(line)
                    if "OK" in line or "ERROR" in line or "CME ERROR" in line:
                        break
            else:
                response_lines.append("TIMEOUT")
            return True, response_lines
        except Exception as e:
            return False, [f"AT命令执行异常: {e}"]
        finally:
            self._close()

    def get_imei(self):
        for attempt in range(2):
            success, response = self.send_command("AT+CGSN=1")
            if success and any("OK" in line for line in response):
                for line in response:
                    match = re.fullmatch(r'\d{15}', line)
                    if "+CGSN:" in line:
                        return line.split(":")[-1].strip()
                    elif match:
                        return line
            time.sleep(0.5)
        return None

    def get_firmware_version(self):
        success, response = self.send_command("ATI")
        if not success or not any("OK" in line for line in response):
            return None
        for line in response:
            if "CTWing" in line or "Rev" in line:
                return line.strip()
        return None

# ==============================================================================
# 2. ExcelReporter 类
# ==============================================================================
class ExcelReporter:
    def __init__(self, filepath="upgrade_test_results.xlsx"):
        self.filepath = filepath
        self.lock = Lock()
        self._initialize_file()

    def _initialize_file(self):
        with self.lock:
            if not os.path.exists(self.filepath):
                workbook = openpyxl.Workbook()
                sheet = workbook.active
                sheet.title = "Test Results"
                headers = ["时间", "IMEI", "初始固件", "最终固件", "测试结果", "端口", "备注"]
                sheet.append(headers)
                workbook.save(self.filepath)
    
    def write_report(self, task):
        with self.lock:
            try:
                workbook = openpyxl.load_workbook(self.filepath)
                sheet = workbook.active
                timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
                data_row = [
                    timestamp,
                    task.imei,
                    task.initial_firmware,
                    task.final_firmware,
                    task.test_result,
                    task.initial_port,
                    task.last_message if task.test_result == "失败" else ""
                ]
                sheet.append(data_row)
                workbook.save(self.filepath)
                return True
            except Exception:
                return False

# ==============================================================================
# 3. LogDialog 类 
# ==============================================================================
class LogDialog(QDialog):
    closed = pyqtSignal()
    def __init__(self, task_id, history, parent=None):
        super().__init__(parent)
        self.setWindowTitle(f"实时日志 - {task_id}")
        self.setGeometry(100, 100, 700, 500)
        self.log_text = QPlainTextEdit(self)
        self.log_text.setReadOnly(True)
        self.log_text.setFont(QFont("Courier New", 10))
        self.log_text.setLineWrapMode(QPlainTextEdit.NoWrap)
        if history:
            self.log_text.setPlainText("\n".join(history))
            self.log_text.verticalScrollBar().setValue(self.log_text.verticalScrollBar().maximum())
        layout = QVBoxLayout(self)
        layout.addWidget(self.log_text)
        self.setLayout(layout)
    def append_log(self, message):
        self.log_text.appendPlainText(message)
    def closeEvent(self, event):
        self.closed.emit()
        super().closeEvent(event)

# ==============================================================================
# 4. UpgradeTask 类 
# ==============================================================================
class UpgradeTask:
    def __init__(self, initial_port):
        self.initial_port = initial_port
        self.imei = "N/A"
        self.initial_firmware = "N/A"
        self.final_firmware = "N/A"
        self.test_result = "未执行"
        self.dl_port = None
        self.status = "未开始"
        self.status_color = QColor('gray')
        self.progress = 0
        self.last_message = f"设备: {initial_port}"
        self.elapsed_time = "00:00"
        self.worker = None
        self.thread = None
        self.is_active = False
        self.log_history = []
        self.log_window = None
    def get_current_display_port(self):
        if self.is_active and self.dl_port:
            return self.dl_port
        return self.initial_port
    def get_all_associated_ports(self):
        ports = {self.initial_port}
        if self.dl_port:
            ports.add(self.dl_port)
        return ports

# ==============================================================================
# 5. DownloadWorker 类 (修改了 __init__ 和 finalize_task)
# ==============================================================================
class DownloadWorker(QObject):
    output_ready = pyqtSignal(str, str)
    progress_updated = pyqtSignal(str, int)
    upgrade_status_updated = pyqtSignal(str, str, QColor)
    upgrade_time_updated = pyqtSignal(str, str)
    dl_port_found = pyqtSignal(str, str)
    finished = pyqtSignal(str)
    requesting_dl_port = pyqtSignal(str)

    def __init__(self, initial_port, image_files, excel_reporter=None): # excel_reporter is now optional
        super().__init__()
        self.initial_port = initial_port
        self.image_files = image_files
        self.excel_reporter = excel_reporter # Can be None
        self.at_communicator = ATCommunicator(initial_port)
        self.current_image_index = 0
        self.is_failed = False
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.send_upgrade_time)
        self.start_time = None
        self.process = QProcess(self)
        self.process.setProcessChannelMode(QProcess.MergedChannels)
        self.process.readyReadStandardOutput.connect(self.handle_stdout)
        self.process.finished.connect(self.process_finished)
        self.stdout_buffer = ""

    def run(self):
        # 只有在需要生成报告时才执行预检查
        if self.excel_reporter:
            success, message = self.get_pre_upgrade_info()
            if not success:
                self.report_failure_and_finish(message)
                return

        if not self.image_files or self.current_image_index >= len(self.image_files):
            self.finalize_upgrade(True)
            return
        
        self.start_time = time.time()
        self.timer.start(1000)
        self.execute_next_command()

    def get_pre_upgrade_info(self):
        task = MyWindow.instance.tasks.get(self.initial_port)
        if not task: return False, "任务对象丢失"
        self.output_ready.emit(self.initial_port, "获取IMEI...")
        task.imei = self.at_communicator.get_imei()
        if not task.imei:
            return False, "获取IMEI失败"
        self.output_ready.emit(self.initial_port, f"IMEI: {task.imei}")
        self.output_ready.emit(self.initial_port, "获取初始固件版本...")
        task.initial_firmware = self.at_communicator.get_firmware_version()
        if not task.initial_firmware:
            return False, "获取初始固件版本失败"
        self.output_ready.emit(self.initial_port, f"初始版本: {task.initial_firmware}")
        return True, "初始信息获取成功"

    def execute_next_command(self):
        if self.current_image_index < len(self.image_files):
            image_file = self.image_files[self.current_image_index]
            exe_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'Bin', 'DownloadShell.exe'))
            command = f'"{exe_path}" -dl {self.initial_port.split("COM")[-1]} "{image_file}"'
            self.process.start(command)
        else:
            self.finalize_upgrade(not self.is_failed)

    def handle_stdout(self):
        data = self.process.readAllStandardOutput().data().decode('utf-8', errors='ignore')
        self.stdout_buffer += data
        while '\n' in self.stdout_buffer:
            line, self.stdout_buffer = self.stdout_buffer.split('\n', 1)
            line = line.strip()
            if not line: continue
            if "Please set download port" in line:
                self.output_ready.emit(self.initial_port, "请求分配DL端口...")
                self.requesting_dl_port.emit(self.initial_port)
            elif not re.match(r'\d+\.\d+%', line):
                self.output_ready.emit(self.initial_port, line)
            if 'fail' in line.lower():
                self.is_failed = True
                self.upgrade_status_updated.emit(self.initial_port, "失败", QColor('red'))
                if self.process.state() == QProcess.Running:
                    self.process.kill()
            self.update_progress(line)

    @pyqtSlot(str)
    def assign_dl_port(self, dl_port_num):
        if dl_port_num:
            self.output_ready.emit(self.initial_port, f"已分配DL端口: COM{dl_port_num}")
            self.process.write(f"{dl_port_num}\n".encode())
        else:
            self.output_ready.emit(self.initial_port, "DL端口分配失败，升级终止。")
            self.is_failed = True
            if self.process.state() == QProcess.Running:
                self.process.kill()

    @pyqtSlot(int, QProcess.ExitStatus)
    def process_finished(self, exitCode, exitStatus):
        if self.is_failed:
            self.finalize_upgrade(False)
            return

        is_normal_exit = (exitStatus == QProcess.NormalExit and exitCode == 0)
        
        # 检查是否还有其他文件需要烧录
        if is_normal_exit and self.current_image_index < len(self.image_files) - 1:
            self.current_image_index += 1
            self.execute_next_command()
        elif is_normal_exit: # 所有文件烧录完成
            self.finalize_upgrade(True)
        else: # 烧录失败
            self.finalize_upgrade(False)

    def report_failure_and_finish(self, message):
        self.timer.stop()
        task = MyWindow.instance.tasks.get(self.initial_port)
        if task and self.excel_reporter:
            task.test_result = "失败"
            task.last_message = message
            self.excel_reporter.write_report(task)
        self.upgrade_status_updated.emit(self.initial_port, "失败", QColor('red'))
        self.output_ready.emit(self.initial_port, message)
        self.finished.emit(self.initial_port)

    def finalize_upgrade(self, success):
        self.timer.stop()
        
        # 只有在需要生成报告时才进行版本验证
        if success and self.excel_reporter:
            self.upgrade_status_updated.emit(self.initial_port, "验证中...", QColor('purple'))
            time.sleep(5) # 等待设备重启并识别AT口
            final_firmware = self.at_communicator.get_firmware_version()
            if final_firmware:
                self.finalize_task(final_firmware)
            else:
                self.report_failure_and_finish("获取最终固件版本失败")
        elif success: # 不需要报告，直接成功
            self.upgrade_status_updated.emit(self.initial_port, "已完成", QColor('green'))
            self.progress_updated.emit(self.initial_port, 100)
            self.finished.emit(self.initial_port)
        else: # 升级失败
            self.report_failure_and_finish("烧录过程失败")

    def finalize_task(self, final_firmware):
        task = MyWindow.instance.tasks.get(self.initial_port)
        if not task: return

        task.final_firmware = final_firmware
        
        if task.initial_firmware and task.final_firmware and task.initial_firmware != task.final_firmware:
            task.test_result = "成功"
            self.upgrade_status_updated.emit(self.initial_port, "已完成", QColor('green'))
        else:
            task.test_result = "失败"
            task.last_message = f"版本未变化或获取失败"
            self.upgrade_status_updated.emit(self.initial_port, "验证失败", QColor('red'))
        
        if self.excel_reporter and not self.excel_reporter.write_report(task):
             self.output_ready.emit(self.initial_port, "写入Excel失败！")

        self.finished.emit(self.initial_port)

    def update_progress(self, output_line):
        total_files = len(self.image_files) if self.image_files else 1
        base_progress = (self.current_image_index / total_files) * 100
        total_steps_per_file = 43
        match = re.search(r'\[(\d+)/(\d+)\] pass\.', output_line)
        if match:
            completed = int(match.group(1))
            file_percentage = (completed / total_steps_per_file)
            total_percentage = int(base_progress + (file_percentage * (100 / total_files)))
            self.progress_updated.emit(self.initial_port, total_percentage)

    def send_upgrade_time(self):
        if self.start_time is not None:
            elapsed_time = time.strftime('%M:%S', time.gmtime(time.time() - self.start_time))
            self.upgrade_time_updated.emit(self.initial_port, elapsed_time)


# ==============================================================================
# 6. MyWindow 类 (多处修改)
# ==============================================================================
class MyWindow(QWidget):
    instance = None
    def __init__(self):
        super().__init__()
        MyWindow.instance = self
        self.tasks = {}
        self.image_files = []
        self.excel_reporter = None # 默认不创建
        self.dl_port_claim_lock = Lock()
        self.claimed_dl_ports = set()
        
        # 新增计数器
        self.active_task_count = 0
        self.active_task_lock = Lock()
        
        self.setup_ui()
        self.apply_styles()
        self.port_timer = QTimer(self)
        self.port_timer.timeout.connect(self.scan_ports)
        self.port_timer.start(1000)
        self.scan_ports()

    def setup_ui(self):
        self.setWindowTitle("XY 智能升级工具")
        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(10, 10, 10, 10)
        main_layout.setSpacing(10)

        top_panel = QHBoxLayout()
        self.select_all_checkbox = QCheckBox("全选")
        self.select_all_checkbox.stateChanged.connect(self.toggle_select_all)
        self.select_all_checkbox.setObjectName("SelectAllCheckbox")

        self.select_image_button = QPushButton(" 选择镜像文件")
        self.select_image_button.clicked.connect(self.select_image_files)
        self.select_image_button.setToolTip("选择一个或多个 .mimgx 镜像文件")
        self.select_image_button.setMinimumHeight(35)

        self.upgrade_button = QPushButton(" 开始升级")
        self.upgrade_button.clicked.connect(self.start_upgrade)
        self.upgrade_button.setToolTip("开始升级所有选中的设备")
        self.upgrade_button.setObjectName("UpgradeButton")
        self.upgrade_button.setMinimumHeight(35)

        self.file_label = QLabel("未选择文件")
        self.file_label.setObjectName("FileLabel")
        
        # 新增：生成报告的复选框
        self.generate_report_checkbox = QCheckBox("生成升级报告")
        self.generate_report_checkbox.setToolTip("勾选后，将在升级前后获取设备信息并记录到Excel")
        
        top_panel.addWidget(self.select_all_checkbox)
        top_panel.addSpacing(10)
        top_panel.addWidget(self.generate_report_checkbox)
        top_panel.addSpacing(20)
        top_panel.addWidget(self.select_image_button)
        top_panel.addWidget(self.file_label, 1)
        top_panel.addStretch()
        top_panel.addWidget(self.upgrade_button)
        main_layout.addLayout(top_panel)

        self.table = QTableWidget()
        self.table.setColumnCount(6)
        self.table.setHorizontalHeaderLabels(["选择", "端口", "升级状态", "升级信息", "进度", "时间"])
        self.table.cellClicked.connect(self.on_cell_clicked)
        header = self.table.horizontalHeader()
        header.setStretchLastSection(False)
        header.setSectionResizeMode(0, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(1, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(2, QHeaderView.ResizeToContents)
        header.setSectionResizeMode(3, QHeaderView.Stretch)
        header.setSectionResizeMode(4, QHeaderView.Interactive)
        header.setSectionResizeMode(5, QHeaderView.ResizeToContents)
        self.table.setColumnWidth(4, 200)
        self.table.verticalHeader().setVisible(False)
        self.table.setSelectionBehavior(QAbstractItemView.SelectRows)
        self.table.setAlternatingRowColors(True)
        self.table.setShowGrid(False)
        main_layout.addWidget(self.table)
        self.setLayout(main_layout)

    # on_worker_finished 方法修改
    def on_worker_finished(self, initial_port):
        with self.dl_port_claim_lock:
            if initial_port in self.tasks:
                task = self.tasks[initial_port]
                if task.dl_port and task.dl_port in self.claimed_dl_ports:
                    self.claimed_dl_ports.remove(task.dl_port)
                task.is_active = False
                task.dl_port = None
                if task.thread:
                    task.thread.quit()
                    task.thread.wait()
                    task.worker.deleteLater()
                    task.thread.deleteLater()
                    task.worker = None
                    task.thread = None

        self.refresh_ui()

        # 使用带锁的计数器来处理完成事件
        with self.active_task_lock:
            self.active_task_count -= 1
            if self.active_task_count == 0:
                self.upgrade_button.setEnabled(True)
                self.select_image_button.setEnabled(True)
                QMessageBox.information(self, "完成", "所有选定任务已结束。")

    # start_upgrade 方法修改
    def start_upgrade(self):
        if not self.image_files:
            QMessageBox.warning(self, "警告", "请先选择镜像文件！")
            return
            
        selected_ports = []
        for row in range(self.table.rowCount()):
            cell_widget = self.table.cellWidget(row, 0)
            if cell_widget:
                checkbox = cell_widget.findChild(QCheckBox)
                if checkbox and checkbox.isChecked() and checkbox.isEnabled():
                    port_item = self.table.item(row, 1)
                    if port_item:
                        port_id = port_item.data(Qt.UserRole)
                        task = self.tasks.get(port_id)
                        if not task or not task.is_active:
                            selected_ports.append(port_id)
                            
        if not selected_ports:
            QMessageBox.warning(self, "警告", "请选择至少一个可用的AT端口。")
            return

        self.upgrade_button.setEnabled(False)
        self.select_image_button.setEnabled(False)
        
        # 根据复选框状态决定是否创建ExcelReporter
        if self.generate_report_checkbox.isChecked():
            self.excel_reporter = ExcelReporter()
        else:
            self.excel_reporter = None

        # 更新计数器
        with self.active_task_lock:
            self.active_task_count = len(selected_ports)

        for port_id in selected_ports:
            task = UpgradeTask(port_id)
            task.is_active = True
            task.status = "准备中"
            task.status_color = QColor('orange')
            self.tasks[port_id] = task
            thread = QThread()
            # 将 excel_reporter 实例或 None 传递给 worker
            worker = DownloadWorker(port_id, self.image_files, self.excel_reporter)
            task.worker = worker
            task.thread = thread

            worker.output_ready.connect(self.on_output_ready)
            worker.progress_updated.connect(self.on_progress_updated)
            worker.upgrade_status_updated.connect(self.on_upgrade_status_updated)
            worker.upgrade_time_updated.connect(self.on_upgrade_time_updated)
            worker.requesting_dl_port.connect(self.on_requesting_dl_port)
            worker.finished.connect(self.on_worker_finished)

            worker.moveToThread(thread)
            thread.started.connect(worker.run)
            thread.start()
        
        self.refresh_ui()

    def apply_styles(self):
        QApplication.setStyle(QStyleFactory.create('Fusion'))
        qss = """
            QWidget {
                background-color: #f0f0f0;
                font-family: "Segoe UI", "Microsoft YaHei", "WenQuanYi Micro Hei", sans-serif;
                font-size: 10pt;
            }
            QTableWidget {
                background-color: #ffffff;
                border: 1px solid #dcdcdc;
                border-radius: 5px;
                gridline-color: #e0e0e0;
                color: #333333;
            }
            QHeaderView::section {
                background-color: #f7f7f7;
                padding: 5px;
                border: none;
                border-bottom: 1px solid #dcdcdc;
                font-weight: bold;
                color: #555;
            }
            QTableWidget::item {
                padding: 5px;
                border: none;
            }
            QTableWidget::item:selected {
                background-color: #e6f7ff;
                color: #000;
            }
            QTableWidget::item:hover {
                background-color: #f5f5f5;
            }
            QTableWidget::alternating-row-color {
                background-color: #fafafa;
            }
            QPushButton {
                background-color: #fdfdfd;
                border: 1px solid #dcdcdc;
                padding: 8px 15px;
                border-radius: 4px;
                color: #333;
            }
            QPushButton:hover {
                background-color: #f0f0f0;
                border-color: #c0c0c0;
            }
            QPushButton:pressed {
                background-color: #e0e0e0;
            }
            QPushButton#UpgradeButton {
                background-color: #28a745;
                color: white;
                font-weight: bold;
                border: none;
            }
            QPushButton#UpgradeButton:hover {
                background-color: #218838;
            }
            QPushButton#UpgradeButton:disabled {
                background-color: #cccccc;
            }
            QLabel#FileLabel {
                color: #666;
                padding-left: 10px;
                font-style: italic;
            }
            QProgressBar {
                border: 1px solid #dcdcdc;
                border-radius: 4px;
                text-align: center;
                background-color: #fdfdfd;
                color: #333;
            }
            QProgressBar::chunk {
                background-color: #28a745;
                border-radius: 3px;
                margin: 1px;
            }
            QMessageBox {
                background-color: #ffffff;
            }
        """
        self.setStyleSheet(qss)

    def on_cell_clicked(self, row, column):
        if column == 3:
            port_item = self.table.item(row, 1)
            if not port_item: return
            task_id = port_item.data(Qt.UserRole)
            task = self.tasks.get(task_id)
            if not task: return
            if task.log_window and task.log_window.isVisible():
                task.log_window.activateWindow()
            else:
                dialog = LogDialog(task_id, task.log_history, self)
                task.log_window = dialog
                dialog.closed.connect(lambda: self.on_log_window_closed(task_id))
                dialog.show()

    def on_log_window_closed(self, task_id):
        task = self.tasks.get(task_id)
        if task:
            task.log_window = None

    def on_output_ready(self, initial_port, message):
        if initial_port in self.tasks:
            task = self.tasks[initial_port]
            task.last_message = message
            task.log_history.append(message)
            if task.log_window:
                task.log_window.append_log(message)
            self.refresh_ui()

    def select_image_files(self, *args):
        files, _ = QFileDialog.getOpenFileNames(self, "选择镜像文件", "", "镜像文件 (*.mimgx)")
        if files:
            self.image_files = files
            self.file_label.setText(f"已选 {len(files)} 个文件: " + os.path.basename(files[0]) + ("..." if len(files) > 1 else ""))
            self.file_label.setToolTip("\n".join(files))
        else:
            self.image_files = []
            self.file_label.setText("未选择文件")
            self.file_label.setToolTip("")

    def toggle_select_all(self, state):
        is_checked = state == Qt.Checked
        for row in range(self.table.rowCount()):
            cell_widget = self.table.cellWidget(row, 0)
            if cell_widget:
                checkbox = cell_widget.findChild(QCheckBox)
                if checkbox and checkbox.isEnabled():
                    checkbox.setChecked(is_checked)

    def add_row(self, port_name, port_description):
        row_position = self.table.rowCount()
        self.table.insertRow(row_position)
        checkbox_container = QWidget()
        chk_layout = QHBoxLayout(checkbox_container)
        checkbox = QCheckBox()
        is_at_port = "AT" in port_description or "USB Serial AT" in port_description
        if not is_at_port:
            checkbox.setEnabled(False)
            checkbox.setToolTip("非AT端口不可选择")
        chk_layout.addWidget(checkbox)
        chk_layout.setAlignment(Qt.AlignCenter)
        chk_layout.setContentsMargins(0, 0, 0, 0)
        self.table.setCellWidget(row_position, 0, checkbox_container)
        port_item = QTableWidgetItem(port_name)
        port_item.setData(Qt.UserRole, port_name)
        self.table.setItem(row_position, 1, port_item)
        status_item = QTableWidgetItem("未开始")
        status_item.setForeground(QBrush(QColor('#888888')))
        self.table.setItem(row_position, 2, status_item)
        info_item = QTableWidgetItem(port_description)
        info_item.setForeground(QBrush(QColor('#666666')))
        self.table.setItem(row_position, 3, info_item)
        progress_bar = QProgressBar()
        progress_bar.setValue(0)
        progress_bar.setTextVisible(True)
        self.table.setCellWidget(row_position, 4, progress_bar)
        self.table.setItem(row_position, 5, QTableWidgetItem("00:00"))

    def scan_ports(self):
        system_ports = {p.device: p.description for p in serial.tools.list_ports.comports()
                        if "XY USB Serial" in p.description}
        table_port_ids = set()
        for row in range(self.table.rowCount()):
            item = self.table.item(row, 1)
            if item:
                table_port_ids.add(item.data(Qt.UserRole))
        for port, desc in system_ports.items():
            if "DL" not in desc and port not in table_port_ids:
                self.add_row(port, desc)
        rows_to_remove = []
        for row in range(self.table.rowCount()):
            item = self.table.item(row, 1)
            if not item: continue
            port_id = item.data(Qt.UserRole)
            task = self.tasks.get(port_id)
            if port_id not in system_ports and (not task or not task.is_active):
                 rows_to_remove.append(row)
        for row in reversed(rows_to_remove):
            self.table.removeRow(row)
        self.refresh_ui()

    def refresh_ui(self):
        for row in range(self.table.rowCount()):
            port_item = self.table.item(row, 1)
            if not port_item: continue
            initial_port_id = port_item.data(Qt.UserRole)
            task = self.tasks.get(initial_port_id)
            if task:
                port_item.setText(task.get_current_display_port())
                if task.is_active and task.dl_port:
                     port_item.setForeground(QBrush(QColor('blue')))
                else:
                     port_item.setForeground(QBrush(QColor('black')))
                self.table.item(row, 2).setText(task.status)
                self.table.item(row, 2).setForeground(QBrush(task.status_color))
                self.table.item(row, 3).setText(task.last_message)
                self.table.cellWidget(row, 4).setValue(task.progress)
                self.table.item(row, 5).setText(task.elapsed_time)

    @pyqtSlot(str)
    def on_requesting_dl_port(self, initial_port):
        task = self.tasks.get(initial_port)
        if not task or not task.worker:
            return
        with self.dl_port_claim_lock:
            max_retries = 5
            for i in range(max_retries):
                available_ports = serial.tools.list_ports.comports()
                found_dl_ports = {
                    p.device: re.search(r'COM(\d+)', p.device) 
                    for p in available_ports 
                    if "USB Serial DL" in p.description
                }
                unclaimed_port_device = None
                for device, match in found_dl_ports.items():
                    if device not in self.claimed_dl_ports:
                        unclaimed_port_device = device
                        break
                if unclaimed_port_device:
                    self.claimed_dl_ports.add(unclaimed_port_device)
                    task.dl_port = unclaimed_port_device
                    port_num = found_dl_ports[unclaimed_port_device].group(1)
                    task.worker.assign_dl_port(port_num)
                    return
                if i < max_retries - 1:
                    time.sleep(1)
            task.worker.assign_dl_port(None)

    def on_progress_updated(self, initial_port, progress):
        if initial_port in self.tasks:
            task = self.tasks[initial_port]
            task.progress = progress
            if progress > 0 and task.status == "准备中":
                task.status = "升级中"
                task.status_color = QColor('blue')
            self.refresh_ui()

    def on_upgrade_status_updated(self, initial_port, status, color):
        if initial_port in self.tasks:
            task = self.tasks[initial_port]
            task.status = status
            task.status_color = color
            self.refresh_ui()

    def on_upgrade_time_updated(self, initial_port, elapsed_time):
        if initial_port in self.tasks:
            task = self.tasks[initial_port]
            if task.is_active:
                task.elapsed_time = elapsed_time
                self.refresh_ui()

    def on_dl_port_found(self, initial_port, dl_port):
        if initial_port in self.tasks:
            self.tasks[initial_port].dl_port = dl_port
            self.refresh_ui()

    def closeEvent(self, event):
        for task in self.tasks.values():
            if task.log_window:
                task.log_window.close()
        if any(t.is_active for t in self.tasks.values()):
             reply = QMessageBox.question(self, '退出', '正在升级，确定要退出吗？',
                                     QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
             if reply == QMessageBox.No:
                 event.ignore()
                 return
        self.port_timer.stop()
        for task in self.tasks.values():
            if task.worker and task.worker.process.state() == QProcess.Running:
                task.worker.process.kill()
            if task.thread and task.thread.isRunning():
                task.thread.quit()
                task.thread.wait()
        event.accept()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MyWindow()
    window.resize(900, 600)
    window.show()
    sys.exit(app.exec_())